# koroFileHeader at Yungoal acer
# Create: 2021-07-30 10:46:28
# LastEdit: 2021-08-02 13:02:12
"""网络安全组，安全规则，查找IP地址
基本利用 Azure Resource Graph 浏览器 实现。
基础查询语句如下：

Resources
| where type == "microsoft.network/networksecuritygroups"
| extend 
	dsr=properties.defaultSecurityRules, 
	sr=properties.securityRules
| mvexpand asr=array_concat(dsr, sr)
| project-away dsr, sr
| extend
	saps=asr.properties.sourceAddressPrefixes,
	sap=asr.properties.sourceAddressPrefix,
	daps=asr.properties.destinationAddressPrefixes,
	dap=asr.properties.destinationAddressPrefix
| mvexpand ip_addr=array_concat(saps, pack_array(sap), daps, pack_array(dap))
| project-away saps, sap, daps, dap
| where ip_addr !in~ ('', '*')
| summarize count() by tostring(ip_addr)
实际使用的更复杂一些。
"""
__author__ = '749B'

import os, sys
BASE_DIR = os.getcwd()
sys.path.append(BASE_DIR)

import argparse
import textwrap
import msrestazure.azure_cloud
# pylint: disable = import-error
from azure_client import Azure_Client, data_format
from utils.blueking import job_start, job_success, job_fail
from utils.tools import remove_empty
from utils.env import init_environs

from typing import Optional, List


def get_query(client: Azure_Client, str_query: str) -> None:
    """查询"""
    # pylint: disable = no-member
    res = client.resource_graph_query_default(
        str_query, 
        [sub.get('subscription_id') for sub in client.get_subscriptions()], 
        )
    # 查询：无结果。
    if not res.data:
        job_success("查询：无结果")
    # 打印结果
    data_format.output_csv(res.data, comma_begin=False)
    print('')

def str_query_base(
    subscription_id: Optional[str] = None,
    resource_group: Optional[str] = None,
    ) -> str:
    strQuery = textwrap.dedent("""\
    Resources
    | where type == "microsoft.network/networksecuritygroups"
    {filter_subscription_id}
    {filter_resource_group}
    | join kind=leftouter (
        Resourcecontainers
        | where type =~ "microsoft.resources/subscriptions"
        | project SubName=name, subscriptionId
    ) on subscriptionId
    | project nsg_name=name, location, resourceGroup, SubName, properties
    | extend 
        dsr=properties.defaultSecurityRules, 
        sr=properties.securityRules
    | mvexpand asr=array_concat(dsr, sr)
    | project-away dsr, sr, properties
    | extend
        saps=asr.properties.sourceAddressPrefixes,
        sap=asr.properties.sourceAddressPrefix,
        daps=asr.properties.destinationAddressPrefixes,
        dap=asr.properties.destinationAddressPrefix
    | mvexpand ip_addr=array_concat(saps, pack_array(sap), daps, pack_array(dap))
    | project-away saps, sap, daps, dap
    """).format(
        filter_subscription_id=f'| where subscriptionId =~ "{subscription_id}"' if subscription_id else '',
        filter_resource_group=f'| where resourceGroup =~ "{resource_group}"' if resource_group else '',
        )
    # 去掉空白行
    strQuery = remove_empty(strQuery) + '\n'
    return strQuery

def str_query_summarize(base_query: str):
    """所有IP地址计数统计"""
    strQuery = base_query + '| summarize count() by tostring(ip_addr)\n'
    return strQuery

def str_query_one_ip(base_query: str, ip_addr: str):
    """单个IP地址的查询
    单个IP有两种表示方法，直接写IP地址，或者带上/32的网络地址。
    所以即使是单个IP地址，也要用 | where ip_addr in~ ("192.168.1.1", "192.168.1.1/32") 来表示
    单个IP地址会生成2种情况，不会进入这个函数，单个网段会进入这个函数。
    """
    strQuery = textwrap.dedent("""\
    {base_query}
    | extend
        priority=asr.properties.priority,
        role_name=asr.name,
        access=asr.properties.access,
        dports=asr.properties.destinationPortRanges,
        dport=asr.properties.destinationPortRange,
        sports=asr.properties.sourcePortRanges,
        sport=asr.properties.sourcePortRange
    | where ip_addr !in~ ("", "*")
    | extend ports=set_difference(array_concat(dports, pack_array(dport), sports, pack_array(sport)), pack_array("*"))
    | project-away dports, dport, sports, sport, asr
    | where ip_addr =~ "{ip_addr}"
    """).format(
        base_query=base_query[:-1],  # 去掉最后一个换行符
        ip_addr=ip_addr
    )
    return strQuery  

def str_query_ips(base_query: str, ips: List[str]):
    """多个IP地址的查询"""
    strQuery = textwrap.dedent("""\
    {base_query}
    | extend
        priority=asr.properties.priority,
        role_name=asr.name,
        access=asr.properties.access,
        dports=asr.properties.destinationPortRanges,
        dport=asr.properties.destinationPortRange,
        sports=asr.properties.sourcePortRanges,
        sport=asr.properties.sourcePortRange
    | where ip_addr in~ ({ips})
    | extend ports=set_difference(array_concat(dports, pack_array(dport), sports, pack_array(sport)), pack_array("*"))
    | project-away dports, dport, sports, sport, asr
    """).format(
        base_query=base_query[:-1],  # 去掉最后一个换行符
        ips=', '.join(['"%s"' % ip for ip in ips])
    )
    return strQuery  

def transform_ip(ip: str) -> List[str]:
    """由于单个IP地址有两种表示方式：只有IP地址，包含网段的/32的地址。
    由于查询语句是精确匹配，只要求用户提供一种形式，然后返回两种形式的列表。
    不考虑是否是合法的IP地址
    """
    if len(ip) < 7:
        return [ip]
    if '/' in ip[-3:-1]:
        # 包含网络地址
        if ip[-3:] == '/32':
            return [ip[:-3], ip]
        else:
            return [ip]
    else:
        return [ip, ip + '/32']

def init_argparse(args: Optional[List[str]] = None) -> argparse.Namespace:
    """初始化：命令行参数"""
    parser = argparse.ArgumentParser(description='导出备份清单')
    
    parser.add_argument('ip_addr', nargs='*', help='匹配的IP地址或网段')
    parser.add_argument('-e', '--envfile', help="存有用户认证信息的env文件的文件名")
    parser.add_argument('-s', '--subscription', help='限定订阅，订阅名称或订阅ID都可')
    parser.add_argument('-g', '--resource_group', help='限定资源组名称')
    parser.add_argument('--show_query', action='store_true', dest='dev_show_query', help='只打印KQL查询语句，检查查询语句问题')

    args = parser.parse_args(args)
    return args

def main(args: Optional[List[str]] = None) -> None:
    """主函数"""
    # 初始化命令行参数
    if args is None:
        # 蓝鲸作业平台有BUG，重做的任务，空着的参数里都会被填充null
        args = [arg for arg in sys.argv[1:] if arg != 'null' and not arg.endswith('=null')]
    args = init_argparse(args)
    if not args.envfile:
        return job_fail("请提供存有用户认证信息的env文件的文件名")
        
    # 初始化环境变量
    envs = init_environs(args.envfile)
    if not envs:
        return job_fail("env文件解密错误")
    if envs['runtime_env'] == 'blueking':
        for key in args.__dict__.keys():
            if key.startswith('dev_'):
                v = getattr(args, key)
                if v:
                    print(f'{key}={v}: 不支持的命令行参数，已经忽略。', file=sys.stderr)
                    setattr(args, key, None)
                    
    # 实例化客户端
    CLOUD = getattr(msrestazure.azure_cloud, envs['cloud_name'])
    client = Azure_Client(CLOUD)

    if args.subscription:
        # pylint: disable = no-member
        subscription = client.get_subscriptionid_by_name(args.subscription)
        if not subscription:
            return job_fail("订阅或订阅ID不存在:%s" % args.subscription)
        args.subscription = subscription

    base_query = str_query_base(args.subscription, args.resource_group)
    str_query = None
    ip_addr_list = []
    args_ip_addr = args.ip_addr if isinstance(args.ip_addr, list) else []
    for ip in args_ip_addr:
        ip_addr_list.extend(transform_ip(ip))
    if len(ip_addr_list) <= 0:
        str_query = str_query_summarize(base_query)
    elif len(ip_addr_list) == 1:
        str_query = str_query_one_ip(base_query, ip_addr_list[0])
    else:
        str_query = str_query_ips(base_query, ip_addr_list)
    assert str_query
    if args.dev_show_query:
        print(str_query)  # 蓝鲸平台上，那么多内容放到下面的job_success里会返回错误。单独print出来。
        return job_success("已打印查询语句，程序结束。")

    get_query(client, str_query)
    if ip_addr_list:
        print("注意：只能精确匹配IP地址，建议再用不指定IP地址来执行一遍，检查是否有包含IP地址的网段。")
    job_success()


if __name__ == '__main__':

    job_start()

    ###### 可在此处开始编写您的脚本逻辑代码
    ###### iJobs中执行脚本成功和失败的标准只取决于脚本最后一条执行语句的返回值
    ###### 如果返回值为0，则认为此脚本执行成功，如果非0，则认为脚本执行失败
    
    main()
